// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <pollux/functions/sparksql/specialforms/at_least_n_non_nulls.h>
#include <pollux/expression/constant_expr.h>
#include <pollux/expression/special_form.h>

using namespace kumo::pollux::exec;

namespace kumo::pollux::functions::sparksql {
    namespace {
        class AtLeastNNonNullsExpr : public SpecialForm {
        public:
            AtLeastNNonNullsExpr(
                TypePtr type,
                bool trackCpuUsage,
                std::vector<ExprPtr> &&inputs,
                int n)
                : SpecialForm(
                      std::move(type),
                      std::move(inputs),
                      AtLeastNNonNullsCallToSpecialForm::kAtLeastNNonNulls,
                      true,
                      trackCpuUsage),
                  n_(n) {
            }

            void evalSpecialForm(
                const SelectivityVector &rows,
                EvalCtx &context,
                VectorPtr &result) override {
                context.ensure_writable(rows, type(), result);
                (*result).clear_nulls(rows);
                auto flatResult = result->as_flat_vector<bool>();
                LocalSelectivityVector activeRowsHolder(context, rows);
                auto activeRows = activeRowsHolder.get();
                POLLUX_DCHECK_NOT_NULL(activeRows);
                auto values = flatResult->mutableValues(rows.end())->asMutable<uint64_t>();
                // If 'n_' <= 0, set result to all true.
                if (n_ <= 0) {
                    bits::orBits(values, rows.asRange().bits(), rows.begin(), rows.end());
                    return;
                }

                bits::andWithNegatedBits(
                    values, rows.asRange().bits(), rows.begin(), rows.end());
                // If 'n_' > inputs_.size() - 1, result should be all false.
                if (n_ > inputs_.size() - 1) {
                    return;
                }

                // Create a temp buffer to track count of non null values for active rows.
                auto nonNullCounts =
                        AlignedBuffer::allocate<int32_t>(activeRows->size(), context.pool(), 0);
                auto *rawNonNullCounts = nonNullCounts->asMutable<int32_t>();
                for (column_index_t i = 1; i < inputs_.size(); ++i) {
                    VectorPtr input;
                    inputs_[i]->eval(*activeRows, context, input);
                    if (context.errors()) {
                        context.deselectErrors(*activeRows);
                    }
                    POLLUX_DYNAMIC_TYPE_DISPATCH_ALL(
                        updateResultTyped,
                        inputs_[i]->type()->kind(),
                        input.get(),
                        n_,
                        context,
                        rawNonNullCounts,
                        flatResult,
                        activeRows);
                    if (activeRows->countSelected() == 0) {
                        break;
                    }
                }
            }

        private:
            void computePropagatesNulls() override {
                propagatesNulls_ = false;
            }

            template<TypeKind Kind>
            void updateResultTyped(
                BaseVector *input,
                int32_t n,
                EvalCtx &context,
                int32_t *rawNonNullCounts,
                FlatVector<bool> *result,
                SelectivityVector *activeRows) {
                using T = typename TypeTraits<Kind>::NativeType;
                exec::LocalDecodedVector decodedVector(context);
                decodedVector.get()->decode(*input, *activeRows);
                bool updateBounds = false;
                activeRows->applyToSelected([&](auto row) {
                    bool nonNull = !decodedVector->is_null_at(row);
                    if constexpr (std::is_same_v<T, double> || std::is_same_v<T, float>) {
                        nonNull = nonNull && !std::isnan(decodedVector->value_at<T>(row));
                    }
                    if (nonNull) {
                        rawNonNullCounts[row]++;
                        if (rawNonNullCounts[row] >= n) {
                            updateBounds = true;
                            result->set(row, true);
                            // Exclude the 'row' from active rows after finding 'n' non-NULL /
                            // non-NaN values.
                            activeRows->setValid(row, false);
                        }
                    }
                });
                if (updateBounds) {
                    activeRows->updateBounds();
                }
            }

            // Result is true if there are at least `n_` non-null and non-NaN values.
            const int n_;
        };
    } // namespace

    TypePtr AtLeastNNonNullsCallToSpecialForm::resolveType(
        const std::vector<TypePtr> & /*argTypes*/) {
        return BOOLEAN();
    }

    ExprPtr AtLeastNNonNullsCallToSpecialForm::constructSpecialForm(
        const TypePtr &type,
        std::vector<ExprPtr> &&compiledChildren,
        bool trackCpuUsage,
        const core::QueryConfig & /*config*/) {
        POLLUX_USER_CHECK_GT(
            compiledChildren.size(),
            1,
            "AtLeastNNonNulls expects to receive at least 2 arguments.");
        POLLUX_USER_CHECK(
            compiledChildren[0]->type()->is_integer(),
            "The first input type should be INTEGER but got {}.",
            compiledChildren[0]->type()->toString());
        auto constantExpr =
                std::dynamic_pointer_cast<exec::ConstantExpr>(compiledChildren[0]);
        POLLUX_USER_CHECK_NOT_NULL(
            constantExpr, "The first parameter should be constant expression.");
        POLLUX_USER_CHECK(
            constantExpr->value()->is_constant_encoding(),
            "The first parameter should be wrapped in constant vector.");
        auto constVector =
                constantExpr->value()->as_unchecked<ConstantVector<int32_t> >();
        POLLUX_USER_CHECK(
            !constVector->is_null_at(0), "The first parameter should not be null.");
        const int32_t n = constVector->value_at(0);
        return std::make_shared<AtLeastNNonNullsExpr>(
            type, trackCpuUsage, std::move(compiledChildren), n);
    }
} // namespace kumo::pollux::functions::sparksql
